# -*- coding: utf-8 -*-
import pymssql
from async_generator import asynccontextmanager
from mg_app_framework import get_store, get_logger
from typing import ContextManager
from .function_utils import FunctionUtils


class _Cursor:
    def __init__(self, cursor: pymssql.Cursor):
        self.cursor = cursor

    async def get_one(self, sql: str, *args, **kwargs):
        return await FunctionUtils.my_run_in_executor(self._get_one, sql, *args, **kwargs)

    async def get_many(self, sql: str, *args, **kwargs):
        return await FunctionUtils.my_run_in_executor(self._get_many, sql, *args, **kwargs)

    def _get_one(self, sql, *args, **kwargs):
        self.execute(sql, *args, **kwargs)
        return self.cursor.fetchone()

    def _get_many(self, sql: str, *args, **kwargs):
        self.execute(sql, *args, **kwargs)
        return self.cursor.fetchall()

    def execute(self, sql, *args, **kwargs):
        try:
            return self.cursor.execute(sql, *args, **kwargs)
        except Exception as e:
            get_logger().error("执行sql出错, sql: {}, args -> {}, kwargs -> {}".format(sql, args, kwargs))
            raise e


class SqlConnection(object):
    def __init__(self, connect_config=None):
        if connect_config is None:
            connect_config = get_store().sql_server_config
        self.connect_config = connect_config
        # pool = PooledDB(
        #     pymssql,
        #     charset="utf8",
        #     **connect_config
        # )
        #
        # self.pool = pool

    @classmethod
    def get_instance(cls):
        if not hasattr(cls, '__connect_config'):
            cls.__connect_config = cls()
        return cls.__connect_config

    @asynccontextmanager
    async def get_cursor(self, as_dict=False, commit=False) -> ContextManager[_Cursor]:
        cursor = None
        conn = None
        try:
            conn = await FunctionUtils.my_run_in_executor(pymssql.connect, **self.connect_config)
            cursor = await FunctionUtils.my_run_in_executor(conn.cursor, as_dict=as_dict)
            yield _Cursor(cursor)
            if commit:
                await FunctionUtils.my_run_in_executor(conn.commit)
        except Exception:
            if conn is not None:
                await FunctionUtils.my_run_in_executor(conn.rollback)
            raise
        finally:
            if cursor is not None:
                await FunctionUtils.my_run_in_executor(cursor.close)
            if conn is not None:
                await FunctionUtils.my_run_in_executor(conn.close)
